# 导入必要的库和模块
from fastapi import APIRouter, Request, HTTPException  # FastAPI框架相关导入，用于创建API路由、处理请求和异常
from fastapi.responses import StreamingResponse  # 用于实现流式响应，实时返回生成结果
from dotenv import load_dotenv  # 用于加载环境变量配置
from app.services.github_service import GitHubService  # 导入GitHub服务，用于获取仓库信息
from app.services.o3_mini_openai_service import OpenAIo3Service  # 导入OpenAI服务，用于生成图表
from app.prompts import (  # 导入提示模板，用于AI模型生成
    SYSTEM_FIRST_PROMPT,  # 第一阶段提示：生成项目架构解释
    SYSTEM_SECOND_PROMPT,  # 第二阶段提示：创建组件映射
    SYSTEM_THIRD_PROMPT,  # 第三阶段提示：生成Mermaid图表
    ADDITIONAL_SYSTEM_INSTRUCTIONS_PROMPT,  # 附加指令提示，用于处理用户自定义指令
    NO_README_INSTRUCTIONS,  # 针对无README仓库的特殊指令
)
from anthropic._exceptions import RateLimitError  # 导入Anthropic API限流异常处理
from pydantic import BaseModel  # 导入Pydantic模型，用于请求体验证和数据校验
from functools import lru_cache  # 导入LRU缓存装饰器，用于缓存函数结果
import re  # 导入正则表达式模块，用于处理Mermaid图表中的点击事件
import json  # 导入JSON处理模块，用于格式化响应数据
import asyncio  # 导入异步IO模块，用于实现异步操作

# 以下是已注释掉的导入，表示曾经使用但现在不再使用的服务
# from app.services.claude_service import ClaudeService  # Claude服务（已弃用，现在使用OpenAI）
# from app.core.limiter import limiter  # 速率限制器（暂时禁用以促进增长）

# 加载环境变量，确保能够访问配置的API密钥等敏感信息
load_dotenv()

print("========== 初始化 OpenAI 生成路由 ==========")

# 创建FastAPI路由器，设置前缀为/openai/mermaid，标签为openai
router = APIRouter(prefix="/openai/mermaid", tags=["openai"])

# 初始化OpenAI o3-mini服务实例，用于后续API调用
# claude_service = ClaudeService()  # Claude服务实例（已注释，不再使用）
o3_service = OpenAIo3Service()  # 初始化OpenAI o3-mini服务实例
print("已初始化 OpenAIo3Service 服务")


# 使用LRU缓存装饰器缓存GitHub数据，避免在成本计算和图表生成时重复API调用
# maxsize=100表示最多缓存100个不同参数组合的结果
@lru_cache(maxsize=100)
def get_cached_github_data(username: str, repo: str, github_pat: str | None = None):
    """获取并缓存GitHub仓库数据
    
    为避免重复API调用，使用缓存机制存储已获取的GitHub仓库数据。
    每次调用时创建新的GitHubService实例，确保使用最新的PAT。
    该缓存数据既用于图表生成，也用于文档生成过程。
    
    Args:
        username (str): GitHub用户名
        repo (str): 仓库名称
        github_pat (str | None): GitHub个人访问令牌（可选）
        
    Returns:
        dict: 包含默认分支、文件树和README的字典
    """
    # 为每次调用创建一个新的GitHubService实例，使用提供的PAT或默认PAT
    current_github_service = GitHubService(pat=github_pat)

    # 获取仓库的默认分支（main或master等）
    default_branch = current_github_service.get_default_branch(username, repo)
    if not default_branch:
        default_branch = "main"  # 如果获取失败，使用"main"作为默认值

    # 获取仓库的完整文件树（所有文件路径的列表）
    file_tree = current_github_service.get_github_file_paths_as_list(username, repo)
    # 获取仓库的README文件内容
    readme = current_github_service.get_github_readme(username, repo)

    # 返回包含所有获取数据的字典
    return {"default_branch": default_branch, "file_tree": file_tree, "readme": readme}


# 定义API请求体模型，使用Pydantic进行数据验证
class ApiRequest(BaseModel):
    """API请求体模型，定义了生成图表所需的参数
    
    Attributes:
        username (str): GitHub用户名
        repo (str): 仓库名称
        instructions (str): 自定义指令（可选）
        api_key (str | None): 用户提供的OpenAI API密钥（可选）
        github_pat (str | None): GitHub个人访问令牌（可选）
    """
    username: str  # GitHub用户名，必填字段
    repo: str  # 仓库名称，必填字段
    instructions: str = ""  # 自定义指令，可选字段，默认为空字符串
    api_key: str | None = None  # 用户提供的OpenAI API密钥，可选字段，默认为None
    github_pat: str | None = None  # GitHub个人访问令牌，可选字段，默认为None


# 定义计算生成成本的路由
@router.post("/cost")
# @limiter.limit("5/minute") # TEMP: 暂时禁用速率限制以促进增长
async def get_generation_cost(request: Request, body: ApiRequest):
    """
    计算生成图表的预估成本
    
    根据仓库文件树和README的令牌数量，计算使用OpenAI API生成图表的预估成本。
    
    Args:
        request (Request): FastAPI请求对象
        body (ApiRequest): 请求体，包含GitHub仓库信息
        
    Returns:
        dict: 包含预估成本或错误信息的字典
    """
    print(f"[OpenAI] 成本计算请求: 用户名={body.username}, 仓库={body.repo}")
    try:
        # 获取缓存的GitHub数据（文件树和README）
        github_data = get_cached_github_data(body.username, body.repo, body.github_pat)
        file_tree = github_data["file_tree"]  # 获取文件树
        readme = github_data["readme"]  # 获取README内容

        # 以下是使用Claude服务计算令牌的代码（已注释）
        # file_tree_tokens = claude_service.count_tokens(file_tree)
        # readme_tokens = claude_service.count_tokens(readme)

        # 使用OpenAI服务计算文件树和README的令牌数量
        file_tree_tokens = o3_service.count_tokens(file_tree)
        readme_tokens = o3_service.count_tokens(readme)

        # 以下是Claude模型的成本计算（已注释）
        # 输入成本：每1M令牌$3（每令牌$0.000003）
        # 输出成本：每1M令牌$15（每令牌$0.000015）
        # input_cost = ((file_tree_tokens * 2 + readme_tokens) + 3000) * 0.000003
        # output_cost = 3500 * 0.000015
        # estimated_cost = input_cost + output_cost

        # OpenAI o3-mini模型的成本计算
        # 输入成本：每1M令牌$1.1（每令牌$0.0000011）
        # 输出成本：每1M令牌$4.4（每令牌$0.0000044）
        input_cost = ((file_tree_tokens * 2 + readme_tokens) + 3000) * 0.0000011  # 计算输入成本
        output_cost = (
            8000 * 0.0000044
        )  # 计算输出成本，8k是基于观察到的输出量（推理过程很消耗令牌）
        estimated_cost = input_cost + output_cost  # 计算总成本

        # 格式化为美元字符串，保留两位小数
        cost_string = f"${estimated_cost:.2f} USD"
        print(f"[OpenAI] 成本计算结果: {cost_string}")
        return {"cost": cost_string}  # 返回成本信息
    except Exception as e:
        print(f"[OpenAI] 成本计算错误: {str(e)}")
        return {"error": str(e)}  # 如果发生异常，返回错误信息


# 处理Mermaid图表中的点击事件，添加GitHub URL
def process_click_events(diagram: str, username: str, repo: str, branch: str) -> str:
    """
    处理Mermaid图表中的点击事件，添加完整的GitHub URL。
    检测路径是文件还是目录，并使用适当的URL格式。
    
    Args:
        diagram (str): Mermaid图表代码
        username (str): GitHub用户名
        repo (str): 仓库名称
        branch (str): 分支名称
        
    Returns:
        str: 处理后的Mermaid图表代码，包含完整的GitHub URL
    """

    def replace_path(match):
        # 从点击事件中提取路径
        path = match.group(2).strip("\"'")

        # 判断路径是文件（有扩展名）还是目录
        is_file = "." in path.split("/")[-1]  # 检查最后一部分是否包含点号，判断是否为文件

        # 构建GitHub URL
        base_url = f"https://github.com/{username}/{repo}"  # GitHub仓库基础URL
        path_type = "blob" if is_file else "tree"  # 文件使用blob路径，目录使用tree路径
        full_url = f"{base_url}/{path_type}/{branch}/{path}"  # 完整的GitHub URL

        # 返回包含完整URL的点击事件
        return f'click {match.group(1)} "{full_url}"'

    # 使用正则表达式匹配点击事件：click 组件名称 "路径/到/某处"
    click_pattern = r'click ([^\s"]+)\s+"([^"]+)"'
    # 使用re.sub替换所有匹配项，将相对路径替换为完整的GitHub URL
    return re.sub(click_pattern, replace_path, diagram)


# 定义流式生成图表的路由
@router.post("/stream")
async def generate_stream(request: Request, body: ApiRequest):
    """
    流式生成系统设计图
    
    这是GitDiagram的核心功能，通过三阶段处理流程生成系统设计图：
    1. 第一阶段：生成项目架构解释
    2. 第二阶段：创建组件映射
    3. 第三阶段：生成Mermaid图表
    
    整个过程使用流式响应，实时返回生成结果。
    支持处理没有README文件的仓库，仅使用文件树结构进行分析。
    
    Args:
        request (Request): FastAPI请求对象
        body (ApiRequest): 请求体，包含GitHub仓库信息和可选的自定义指令
        
    Returns:
        StreamingResponse: 流式响应对象，包含生成过程和最终结果
    """
    print(f"[OpenAI] 生成请求: 用户名={body.username}, 仓库={body.repo}, 是否有API密钥={bool(body.api_key)}")

    # 验证指令长度，避免过长的指令
    if len(body.instructions) > 1000:
        print("[OpenAI] 错误: 指令长度超过限制")
        return {"error": "Instructions exceed maximum length of 1000 characters"}

    try:
        # 流式响应的事件生成器
        async def event_generator():
            """
            生成事件流的异步生成器函数
            
            在函数内部包含三个阶段:
            1. 生成项目架构解释
            2. 创建组件映射
            3. 生成Mermaid图表
            
            每个阶段的结果都通过事件流实时返回给客户端
            
            Yields:
                str: 格式化为Server-Sent Events的JSON数据
            """
            
            try:
                # 获取缓存的GitHub数据（默认分支、文件树和README）
                github_data = get_cached_github_data(
                    body.username, body.repo, body.github_pat
                )
                default_branch = github_data["default_branch"]  # 获取默认分支
                file_tree = github_data["file_tree"]  # 获取文件树
                readme = github_data["readme"]  # 获取README内容（可能为空字符串）
                
                # 标记是否存在README
                has_readme = len(readme.strip()) > 0
                if not has_readme:
                    print(f"[OpenAI] 注意: 仓库 {body.username}/{body.repo} 没有README文件，将仅使用文件树进行分析")
                    # 确保README变量不是None，防止后续计算token出错
                    readme = ""
                
                print(f"[OpenAI] 已获取GitHub数据: 默认分支={default_branch}, 文件树长度={len(file_tree)}, 有README={has_readme}")

                # 发送初始状态，通知客户端生成过程已开始
                yield f"data: {json.dumps({'status': 'started', 'message': 'Starting generation process...'})}\n\n"
                
                await asyncio.sleep(0.1)  # 短暂延迟，确保客户端收到初始状态
                
                # 令牌数量检查，确保不超过模型的上下文长度限制
                combined_content = f"{file_tree}\n{readme}"  # 合并文件树和README内容
                token_count = o3_service.count_tokens(combined_content)  # 计算令牌数量
                print(f"[OpenAI] 令牌数量检查: 总数={token_count}")

                # 如果令牌数量超过限制且未提供自定义API密钥，返回错误
                if 50000 < token_count < 195000 and not body.api_key:
                    yield f"data: {json.dumps({'error': f'File tree and README combined exceeds token limit (50,000). Current size: {token_count} tokens. This GitHub repository is too large for my wallet, but you can continue by providing your own OpenAI API key.'})}\n\n"
                    return
                # 如果令牌数量超过模型最大上下文长度，返回错误
                elif token_count > 195000:
                    yield f"data: {json.dumps({'error': f'Repository is too large (>195k tokens) for analysis. OpenAI o3-mini\'s max context length is 200k tokens. Current size: {token_count} tokens.'})}\n\n"
                    return

                # 准备提示模板
                # 如果有自定义指令，则将其添加到第一阶段和第三阶段提示中
                first_system_prompt = SYSTEM_FIRST_PROMPT  # 第一阶段提示模板
                third_system_prompt = SYSTEM_THIRD_PROMPT  # 第三阶段提示模板
                
                # 如果没有README，添加特殊指令
                if not has_readme:
                    first_system_prompt = first_system_prompt + "\n" + NO_README_INSTRUCTIONS
                
                if body.instructions:
                    # 将自定义指令添加到第一阶段提示
                    first_system_prompt = (
                        first_system_prompt
                        + "\n"
                        + ADDITIONAL_SYSTEM_INSTRUCTIONS_PROMPT
                    )
                    # 将自定义指令添加到第三阶段提示
                    third_system_prompt = (
                        third_system_prompt
                        + "\n"
                        + ADDITIONAL_SYSTEM_INSTRUCTIONS_PROMPT
                    )

                # 第一阶段：获取项目架构解释
                print("[OpenAI] 开始第一阶段: 生成项目架构解释")
                explanation = ""  # 初始化解释文本
                # 调用OpenAI API流式生成项目架构解释
                async for chunk in o3_service.call_o3_api_stream(
                    system_prompt=first_system_prompt,  # 使用第一阶段提示模板
                    data={
                        "file_tree": file_tree,  # 提供文件树
                        "readme": readme,  # 提供README
                        "instructions": body.instructions,  # 提供自定义指令
                    },
                    api_key=body.api_key,  # 使用用户提供的API密钥（如果有）
                    reasoning_effort="medium",  # 使用中等推理努力级别
                ):
                    explanation += chunk  # 累积解释文本
                    # 流式返回解释块给客户端
                    yield f"data: {json.dumps({'status': 'explanation_chunk', 'chunk': chunk})}\n\n"
                
                print("[OpenAI] 第一阶段完成: 解释长度=" + str(len(explanation)))

                # 检查是否包含BAD_INSTRUCTIONS标记，表示指令无效
                if "BAD_INSTRUCTIONS" in explanation:
                    yield f"data: {json.dumps({'error': 'Invalid or unclear instructions provided'})}\n\n"
                    return

                # 第二阶段：创建组件映射
                print("[OpenAI] 开始第二阶段: 创建组件映射")
                full_second_response = ""  # 初始化组件映射响应
                # 调用OpenAI API流式生成组件映射
                async for chunk in o3_service.call_o3_api_stream(
                    system_prompt=SYSTEM_SECOND_PROMPT,  # 使用第二阶段提示模板
                    data={"explanation": explanation, "file_tree": file_tree},  # 提供解释和文件树
                    api_key=body.api_key,  # 使用用户提供的API密钥（如果有）
                    reasoning_effort="low",  # 使用低推理努力级别
                ):
                    full_second_response += chunk  # 累积组件映射响应
                    # 流式返回映射块给客户端
                    yield f"data: {json.dumps({'status': 'mapping_chunk', 'chunk': chunk})}\n\n"
                
                print("[OpenAI] 第二阶段完成: 组件映射长度=" + str(len(full_second_response)))

                # 提取组件映射（从标签中提取）
                # 注释表明这部分可能不再需要，但暂时保留
                start_tag = "<component_mapping>"  # 组件映射开始标签
                end_tag = "</component_mapping>"  # 组件映射结束标签
                # 从响应中提取组件映射文本
                component_mapping_text = full_second_response[
                    full_second_response.find(start_tag) : full_second_response.find(
                        end_tag
                    )
                ]

                # 第三阶段：生成Mermaid图表
                print("[OpenAI] 开始第三阶段: 生成Mermaid图表")
                mermaid_code = ""  # 初始化Mermaid代码
                # 调用OpenAI API流式生成Mermaid图表
                async for chunk in o3_service.call_o3_api_stream(
                    system_prompt=third_system_prompt,  # 使用第三阶段提示模板
                    data={
                        "explanation": explanation,  # 提供解释
                        "component_mapping": component_mapping_text,  # 提供组件映射
                        "instructions": body.instructions,  # 提供自定义指令
                    },
                    api_key=body.api_key,  # 使用用户提供的API密钥（如果有）
                    reasoning_effort="medium",  # 使用中等推理努力级别
                ):
                    mermaid_code += chunk  # 累积Mermaid代码
                    # 流式返回图表块给客户端
                    yield f"data: {json.dumps({'status': 'diagram_chunk', 'chunk': chunk})}\n\n"
                
                print("[OpenAI] 第三阶段完成: Mermaid代码长度=" + str(len(mermaid_code)))

                # 处理最终图表：移除Markdown代码块标记
                mermaid_code = mermaid_code.replace("```mermaid", "").replace("```", "")
                # 检查是否包含BAD_INSTRUCTIONS标记，表示指令无效
                if "BAD_INSTRUCTIONS" in mermaid_code:
                    print("[OpenAI] 错误: 无效或不明确的指令")
                    yield f"data: {json.dumps({'error': 'Invalid or unclear instructions provided'})}\n\n"
                    return

                # 处理图表中的点击事件，添加完整的GitHub URL
                processed_diagram = process_click_events(
                    mermaid_code, body.username, body.repo, default_branch
                )
                
                print("[OpenAI] 生成完成: 返回最终结果")

                # 发送最终结果给客户端
                yield f"data: {json.dumps({
                    'status': 'complete',  # 状态：完成
                    'diagram': processed_diagram,  # 处理后的图表
                    'explanation': explanation,  # 项目架构解释
                    'mapping': component_mapping_text  # 组件映射
                })}\n\n"

            except Exception as e:
                # 如果发生异常，返回错误信息
                print(f"[OpenAI] 生成过程中出错: {str(e)}")
                yield f"data: {json.dumps({'error': str(e)})}\n\n"

        # 返回流式响应对象
        return StreamingResponse(
            event_generator(),  # 使用事件生成器
            media_type="text/event-stream",  # 设置媒体类型为事件流
            headers={
                "X-Accel-Buffering": "no",  # 提示Nginx不要缓冲响应
                "Cache-Control": "no-cache",  # 禁用缓存
                "Connection": "keep-alive",  # 保持连接
            },
        )
    except Exception as e:
        # 如果发生异常，返回错误信息
        print(f"[OpenAI] 请求处理出错: {str(e)}")
        return {"error": str(e)}
